import torch
import torchvision
from torch.utils import data
from torchvision import transforms
from d2l import torch as d2l
d2l.use_svg_display()

def download_data(resize=None):
    trans = transforms.ToTensor()
    if resize:
        trans.insert(0,transforms.Resize(resize))
        trans = transforms.Compose(trans)
    mnist_train = torchvision.datasets.FashionMNIST(root='DL\\SoftmaxRegression\\data',train=True,transform=trans,download=True)
    mnist_test = torchvision.datasets.FashionMNIST(root='DL\\SoftmaxRegression\\data',train=False,transform=trans,download=True)
    return mnist_train,mnist_test

def load_data_fashion_mnist(batch_size,num_workers,resize=None):
    mnist_train,mnist_test = download_data(resize)
    return (data.DataLoader(mnist_train,batch_size,shuffle=True,num_workers=num_workers)),\
            (data.DataLoader(mnist_test,batch_size,shuffle=True,num_workers=num_workers))

######################################################################
mnist_train,mnist_test = download_data()
(width,height) = mnist_train[0][0].shape[1:]
#数据显示测试
def get_labels(labels):
    text_labels = mnist_train.classes
    return [text_labels[i] for i in labels] 

def show_images(imgs,r,c,titles=None,scale=1.5):
    figSize = (c*scale,r*scale)
    _,axes = d2l.plt.subplots(r,c,figsize=figSize)
    axes = axes.flatten()
    for i,(ax,img) in enumerate(zip(axes,imgs)):
        if torch.is_tensor(img):
            ax.imshow(img.numpy())
        else:
            ax.imshow(img)
        ax.axes.get_xaxis().set_visible(False)
        ax.axes.get_yaxis().set_visible(False)
        if titles:
            ax.set_title(titles[i])
    d2l.plt.show()
batch_size = 18
X,y = next(iter(data.DataLoader(mnist_train,batch_size=batch_size)))
show_images(X.reshape(batch_size,width,height),3,6,titles=get_labels(y))

#小批量测速
b_s = 256
num_workers = 4
train_iter = data.DataLoader(mnist_train,b_s,shuffle=True,num_workers=num_workers)
timer = d2l.Timer()
for X,y in train_iter:
    continue
print('%.2f'%timer.stop())